# 绘制数据集的ground truth轨迹 （目前支持kitti， 4seasons, urban）

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D


def draw_kitti(filename):
    dataset = []
    groundtruth = []
    with open(filename) as f:
        list_file = f.readlines()

        # 将每一行数据转为数组
        for i in range(len(list_file)):
            list_line = list_file[i].split(' ')
            # 将元素由字符串转为float
            list_line = list(map(float, list_line))
            # 向量转矩阵
            list_line = np.array(list_line)
            list_line.resize(3, 4)
            dataset.append(list_line)
            groundtruth.append([list_line[0, 3], list_line[2, 3]])
        # 最后得到两个numpu矩阵，dataset是存放所有真值的矩阵，groundtruth是存放xy真值的矩阵
        dataset = np.array(dataset)
        groundtruth = np.array(groundtruth)

    x_data = []
    y_data = []
    for i in range(len(dataset)):
        x_data.append(float(dataset[i][0, 3]))
        y_data.append(float(dataset[i][2, 3]))

    print(x_data)
    # 绘制
    plt.plot(x_data, y_data)
    plt.show()


def draw_tum(filename,datasetname="urban"):
    dataset = np.loadtxt(filename, delimiter=' ')
    print(dataset.shape)

    x_data = dataset[:, 1] - dataset[0, 1]
    y_data = dataset[:, 2] - dataset[0, 2]
    z_data = dataset[:, 3] - dataset[0, 3]

    # 创建一个三维坐标系
    fig = plt.figure()
    ax = fig.add_subplot(121, projection='3d')
    ax1 = fig.add_subplot(122)
    # 绘制散点图
    ax.scatter(x_data, y_data, z_data, marker='o', s=1)
    # 绘制折线图
    if datasetname=="urban":
        ax1.plot(x_data, y_data, color='r')
    elif datasetname=="kitti":
        ax1.plot(x_data, z_data, color='r')
    # 设置坐标轴标签
    ax.set_xlabel('X Label')
    ax.set_ylabel('Y Label')
    ax.set_zlabel('Z Label')

    # 设置XYZ坐标轴刻度一致
    max_range = max(max(abs(x_data)), max(abs(y_data)), max(abs(z_data)))
    # ax.set_xlim([-max_range, max_range])
    # ax.set_ylim([-max_range, max_range])
    ax.set_zlim([-max_range, max_range])

    plt.show()


def draw_4Seasons(filename):
    dataset = []
    groundtruth = []
    with open(filename) as f:
        list_file = f.readlines()

        # 将每一行数据转为数组
        for i in range(len(list_file)):
            list_line = list_file[i].split()
            # 将元素由字符串转为float
            list_line = list(map(float, list_line))
            # 向量转矩阵
            list_line = np.array(list_line)
            dataset.append(list_line)
            groundtruth.append([list_line[1], list_line[2], list_line[3]])
        # 最后得到两个numpu矩阵，dataset是存放所有真值的矩阵，groundtruth是存放xy真值的矩阵
        dataset = np.array(dataset)
        groundtruth = np.array(groundtruth)

    x_data = []
    y_data = []
    z_data = []
    for i in range(len(dataset)):
        x_data.append(float(groundtruth[i][0]))
        y_data.append(float(groundtruth[i][1]))
        z_data.append(float(groundtruth[i][2]))

    # 创建一个三维坐标系
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')
    # 绘制
    # 绘制散点图
    ax.scatter(x_data, y_data, z_data, marker='o', s=1)
    # 设置坐标轴标签
    ax.set_xlabel('X Label')
    ax.set_ylabel('Y Label')
    ax.set_zlabel('Z Label')
    plt.show()

def draw_urban_vrs_gps(filename):
    # 读取VRS GPS数据（ground truth）
    # format: [timestamp, latitude, longitude, x coordinate, y coordinate, altitude, fix state, number of satellite,
    # horizontal precision, latitude std, longitude std, altitude std, heading validate flag , magnetic global heading,
    # speed in knot, speed in km, GNVTG mode, ortometric altitude]
    # [时间戳，纬度，经度，x坐标，y坐标，高度，固定状态，卫星数量，
    # 水平精度，纬度std，经度std，海拔std，航向确认标志，磁性全局航向，
    # 结速， 公里，GNVTG模式，高度计高度]
    with open(filename, encoding='utf-8') as file:
        vrs_gps_data = np.loadtxt(file, str, delimiter=",")
        print(vrs_gps_data.shape)
    x_init, y_init = vrs_gps_data[0][3], vrs_gps_data[0][4]
    x_init = float(x_init)
    y_init = float(y_init)

    # 转换其它经纬度坐标并计算差值
    x_data = [0.0]
    y_data = [0.0]
    for i in range(len(vrs_gps_data)):
        tmp_x, tmp_y = float(vrs_gps_data[i][3]), float(vrs_gps_data[i][4])
        x_data.append(tmp_x - x_init)
        y_data.append(tmp_y - y_init)

    print(x_data)
    # 绘制
    plt.plot(x_data, y_data)
    plt.show()


if __name__ == '__main__':
    # kitti
    # filename = '/media/daybeha/Elements/SLAM_dataset/Kitti/sequences/poses/00.txt'
    # draw_kitti(filename)

    # # 4seasons
    # filename = '/media/daybeha/Elements/SLAM_dataset/4Seasons/recording_2020-03-24_17-36-22/result.txt'
    # draw_4Seasons(filename)

    # # urban
    # filename = "/media/daybeha/LTFM2/dataset/urban/urban39-pankyo/sensor_data/vrs_gps.csv"
    # draw_urban_vrs_gps(filename)

    # tum
    filename = "/media/daybeha/LTFM2/dataset/urban/urban38-pankyo/global_pose_tum.txt"
    # filename = "/home/daybeha/Documents/github/ORB_SLAM3_detailed_comments-master/Examples/src/ORB_SLAM3/results/trajectory_tum.txt"
    draw_tum(filename)
